{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from cot import Collection\n",
    "from cot.generate import FRAGMENTS\n",
    "from rich.pretty import pprint\n",
    "import json\n",
    "from sprint_utils import model_sprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all datasets of thoughtsource_100:\n",
    "\n",
    "# \"commonsense_qa\", \"validation\"\n",
    "# \"med_qa\", \"test\"\n",
    "# \"medmc_qa\", \"validation\"\n",
    "# \"open_book_qa\", \"test\"\n",
    "# \"strategy_qa\", \"train\"\n",
    "# \"worldtree\", \"test\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# None,         # no cot-trigger\n",
    "# 'kojima-01',  # Answer: Let's think step by step.\n",
    "# 'kojima-02',  # Answer: First,\n",
    "# 'kojima-03',  # Answer: Let's think about this logically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MOCK_API\n",
    "number_examples = 10\n",
    "dataset_name_splits = [(\"commonsense_qa\", \"validation\")] \n",
    "api_service_model = [(\"mock_api\", \"mock_model\", 1)]\n",
    "cot_trigger_keys = [\"kojima-01\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading commonsense_qa...\n",
      "Evaluating commonsense_qa validation...\n"
     ]
    }
   ],
   "source": [
    "model_sprint(dataset_name_splits, api_service_model, cot_trigger_keys, number_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# real API\n",
    "number_examples = 100\n",
    "dataset_name_splits = [(\"commonsense_qa\", \"validation\")] \n",
    "api_service_model = [(\"openai\", \"text-davinci-003\", 1)]\n",
    "cot_trigger_keys = [\"kojima-01\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading commonsense_qa...\n",
      "Evaluating commonsense_qa validation...\n"
     ]
    }
   ],
   "source": [
    "model_sprint(dataset_name_splits, api_service_model, cot_trigger_keys, number_examples)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ccbfa654f25866afe66a1c016d0b518a994fbe20a93c2d8b432dbe114020e2d2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
